/**
 * @author Gao Jian
 * @email gaoj@njupt.edu.cn
 * @create date 2021-06-06
 * @desc 第2步
 */

#include "defs.h"

__global__ void init_regions(unsigned int *regions, unsigned char *data,
                             unsigned char value, unsigned int nw,
                             unsigned int nh);
__global__ void flip_regions(unsigned char *active, unsigned int *regions,
                             unsigned char mode, unsigned int nw,
                             unsigned int nh);
__global__ void area_regions(unsigned int *areas, unsigned int *regions,
                             unsigned int nw, unsigned int nh);
__global__ void kcleanparts(unsigned char *data, unsigned int *areas,
                            unsigned int *regions, unsigned char value,
                            unsigned int threshold, unsigned int nw,
                            unsigned int nh);

__global__ void kbua2a(unsigned *nitm, unsigned *nsum, unsigned char *data,
                       unsigned resolution, unsigned int nw, unsigned int nh) {
  unsigned nz = nw * nh;
  unsigned increment = blockDim.x * gridDim.x;
  unsigned nits = nz / increment + 1;
  unsigned idx = blockDim.x * blockIdx.x + threadIdx.x;

  for (unsigned i = 0; i < nits; i++, idx += increment) {
    if (idx >= nz)
      break;

    if (data[idx] != AR_BUA) {
      continue;
    }

    nitm[idx] = 0;
    nsum[idx] = 0;

    int iw = idx % nw;
    int ih = idx / nw;

    float radius = 564. / resolution;

    int ws[2] = {iw - int(radius) <= 0 ? 0 : iw - int(radius),
                 iw + int(radius) >= nw - 1 ? nw - 1 : iw + int(radius)};
    int hs[2] = {ih - int(radius) <= 0 ? 0 : ih - int(radius),
                 ih + int(radius) >= nh - 1 ? nh - 1 : ih + int(radius)};

    float radius2 = radius * radius;

    for (int h = hs[0]; h <= hs[1]; h++) {
      for (int w = ws[0]; w <= ws[1]; w++) {
        int dw = w - iw;
        int dh = h - ih;
        if (dw * dw + dh * dh > radius2)
          continue;
        nsum[idx]++;
        if (data[nw * h + w] == AR_BUA)
          nitm[idx]++;
      }
    }
  }
}

__global__ void kbua2b(unsigned char *output, unsigned char *input,
                       unsigned *nitm, unsigned *nsum, unsigned int nw,
                       unsigned int nh) {
  unsigned nz = nw * nh;
  unsigned increment = blockDim.x * gridDim.x;
  unsigned nits = nz / increment + 1;
  unsigned idx = blockDim.x * blockIdx.x + threadIdx.x;

  for (unsigned i = 0; i < nits; i++, idx += increment) {
    if (idx >= nz)
      break;

    if (input[idx] != AR_BUA || nitm[idx] < nsum[idx] * 0.25)
      output[idx] = AR_NUL;
    else
      output[idx] = AR_BUA;
  }
}

__global__ void kbua3a(unsigned char *output, unsigned char *input,
                       unsigned resolution, unsigned int nw, unsigned int nh) {
  unsigned nz = nw * nh;
  unsigned increment = blockDim.x * gridDim.x;
  unsigned nits = nz / increment + 1;
  unsigned idx = blockDim.x * blockIdx.x + threadIdx.x;

  for (unsigned i = 0; i < nits; i++, idx += increment) {
    if (idx >= nz)
      break;

    output[idx] = input[idx];

    if (input[idx] == AR_BUA) {
      continue;
    }

    int iw = idx % nw;
    int ih = idx / nw;

    float radius = 100. / resolution;

    int ws[2] = {iw - int(radius) <= 0 ? 0 : iw - int(radius),
                 iw + int(radius) >= nw - 1 ? nw - 1 : iw + int(radius)};
    int hs[2] = {ih - int(radius) <= 0 ? 0 : ih - int(radius),
                 ih + int(radius) >= nh - 1 ? nh - 1 : ih + int(radius)};

    float radius2 = radius * radius;

    for (int h = hs[0]; h <= hs[1]; h++) {
      for (int w = ws[0]; w <= ws[1]; w++) {
        int dw = w - iw;
        int dh = h - ih;
        if (dw * dw + dh * dh > radius2)
          continue;
        if (input[nw * h + w] == AR_BUA) {
          output[idx] = AR_BUA;
          break;
        }
      }
    }
  }
}

int label_bua(unsigned *temp, unsigned char *output, unsigned char *input,
              unsigned resolution, unsigned int nw, unsigned int nh) {
  unsigned int nz = nw * nh;
  unsigned char *d_input = NULL;
  unsigned char *d_output = NULL;
  unsigned *d_item = NULL;
  unsigned *d_nsum = NULL;
  unsigned char *tmp = NULL;
  unsigned *d_regions = NULL;
  unsigned *d_areas = NULL;

  cudaMalloc((void **)&d_input, nz * sizeof(unsigned char));
  cudaMalloc((void **)&d_output, nz * sizeof(unsigned char));
  cudaMalloc((void **)&d_item, nz * sizeof(unsigned));
  cudaMalloc((void **)&d_nsum, nz * sizeof(unsigned));
  cudaMemcpy(d_input, input, nz * sizeof(unsigned char),
             cudaMemcpyHostToDevice);

  kbua2a<<<NUM_BLOCKS, NUM_THREADS>>>(d_item, d_nsum, d_input, resolution, nw,
                                      nh);
  kbua2b<<<NUM_BLOCKS, NUM_THREADS>>>(d_output, d_input, d_item, d_nsum, nw,
                                      nh);
  tmp = (unsigned char *)d_nsum;
  kbua3a<<<NUM_BLOCKS, NUM_THREADS>>>(tmp, d_output, resolution, nw, nh);
  cudaMemcpy(d_output, tmp, nz * sizeof(unsigned char),
             cudaMemcpyDeviceToDevice);

  d_regions = d_item;
  d_areas = d_nsum;
  init_regions<<<NUM_BLOCKS, NUM_THREADS>>>(d_regions, d_output, AR_NUL, nw,
                                            nh);

  unsigned char *ctmp = (unsigned char *)d_areas;

  unsigned char active;
  unsigned nit = 0;
  do {
    active = 0;
    cudaMemcpy(ctmp, &active, sizeof(unsigned char), cudaMemcpyHostToDevice);
    flip_regions<<<NUM_BLOCKS, NUM_THREADS>>>(ctmp, d_regions, 0, nw, nh);
    flip_regions<<<NUM_BLOCKS, NUM_THREADS>>>(ctmp, d_regions, 2, nw, nh);
    flip_regions<<<NUM_BLOCKS, NUM_THREADS>>>(ctmp, d_regions, 1, nw, nh);
    flip_regions<<<NUM_BLOCKS, NUM_THREADS>>>(ctmp, d_regions, 3, nw, nh);
    cudaMemcpy(&active, ctmp, sizeof(unsigned char), cudaMemcpyDeviceToHost);
    nit++;
    printf("+");
    if (nit % 50 == 0)
      printf("  [%3d]\n  ", nit);
  } while (active);
  printf("  [%3d]\n", nit);

  nppsSet_32u(0, d_areas, nz);
  unsigned int threshold = (unsigned int)(2000000 / resolution / resolution);

  area_regions<<<NUM_BLOCKS, NUM_THREADS>>>(d_areas, d_regions, nw, nh);
  kcleanparts<<<NUM_BLOCKS, NUM_THREADS>>>(d_output, d_areas, d_regions, AR_BUA,
                                           threshold, nw, nh);

  cudaMemcpy(output, d_output, nz * sizeof(unsigned char),
             cudaMemcpyDeviceToHost);
  if (temp)
    cudaMemcpy(temp, d_item, nz * sizeof(unsigned), cudaMemcpyDeviceToHost);
  cudaFree(d_nsum);
  cudaFree(d_item);
  cudaFree(d_output);
  cudaFree(d_input);

  return 0;
}